import sys
import os
import time
from ctypes import addressof, c_ubyte

import gxipy as gx
from gxipy.gxidef import GxPixelFormatEntry, DxValidBit
from gxipy.ImageFormatConvert import ImageFormatConvert

import numpy as np
import cv2

from PyQt5.QtCore import Qt, QTimer, QThread, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import (
    QApplication, QWidget, QLabel, QPushButton, QLineEdit,
    QFileDialog, QHBoxLayout, QVBoxLayout, QMessageBox
)

# ------------------------- 工业相机线程 -------------------------
class CameraWorker(QThread):
    frame_ready = pyqtSignal(np.ndarray)  # 用于 GUI 显示

    def __init__(self, cam_index=1, parent=None):
        super().__init__(parent)
        self.cam_index = cam_index
        self.running = False
        self.recording = False
        self.video_writer = None
        self.log_file = None

    def run(self):
        # 初始化相机
        dm = gx.DeviceManager()
        num, _ = dm.update_all_device_list()
        if num < self.cam_index:
            return
        self.cam = dm.open_device_by_index(self.cam_index)
        feat = self.cam.get_remote_device_feature_control()
        # 曝光 & 软触发
        feat.get_float_feature("ExposureTime").set(70000.0)
        feat.get_enum_feature("TriggerMode").set("On")
        feat.get_enum_feature("TriggerSource").set("Software")
        # 转换器
        self.img_conv = dm.create_image_format_convert()
        # 注册回调
        self.cam.data_stream[0].register_capture_callback(self._on_frame)
        # 获取触发命令
        self.trig = feat.get_command_feature("TriggerSoftware")
        # 开始采集
        self.cam.stream_on()
        self.running = True
        # 主循环：仅保持线程存活
        while self.running:
            time.sleep(0.01)
        # 清理
        self.cam.stream_off()
        self.cam.data_stream[0].unregister_capture_callback()
        self.cam.close_device()
        if self.recording:
            self._stop_recording()

    def _on_frame(self, raw):
        """回调：转换 & 显示 & 录制"""
        # 转换到 RGB8 → BGR
        fmt = GxPixelFormatEntry.RGB8
        self.img_conv.set_dest_format(fmt)
        self.img_conv.set_valid_bits(get_best_valid_bits(raw.get_pixel_format()))
        size = self.img_conv.get_buffer_size_for_conversion(raw)
        buf = (c_ubyte * size)()
        self.img_conv.convert(raw, addressof(buf), size, False)
        arr = np.frombuffer(buf, dtype=np.uint8, count=size)
        h, w = raw.get_height(), raw.get_width()
        rgb = arr.reshape(h, w, 3)
        bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

        # 如果在录制，写 Video + Log
        if self.recording:
            self.video_writer.write(bgr)
            ts = time.time()
            self.log_file.write(f"{self.frame_idx}\t{ts:.6f}\n")
            self.frame_idx += 1

        # 发射信号给 GUI 刷新
        self.frame_ready.emit(bgr)

    def trigger_once(self):
        """触发一次软件采集"""
        if self.running:
            self.trig.send_command()

    def start_recording(self, video_path, log_path):
        """开始录制：打开 VideoWriter 和日志"""
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        # 等待第一帧来获得尺寸
        self._wait_first_frame()
        h, w = self.first_frame.shape[:2]
        self.video_writer = cv2.VideoWriter(video_path, fourcc, 30, (w, h))
        self.log_file = open(log_path, "w", encoding="utf-8")
        self.log_file.write("frame_idx\ttimestamp\n")
        self.frame_idx = 0
        self.recording = True

    def _wait_first_frame(self):
        """内部：临时回调接收第一帧"""
        self.first_frame = None
        def one_time_cb(raw):
            # 复制逻辑：仅取第一帧
            fmt = GxPixelFormatEntry.RGB8
            self.img_conv.set_dest_format(fmt)
            self.img_conv.set_valid_bits(get_best_valid_bits(raw.get_pixel_format()))
            size = self.img_conv.get_buffer_size_for_conversion(raw)
            buf = (c_ubyte * size)()
            self.img_conv.convert(raw, addressof(buf), size, False)
            arr = np.frombuffer(buf, dtype=np.uint8, count=size)
            h, w = raw.get_height(), raw.get_width()
            rgb = arr.reshape(h, w, 3)
            self.first_frame = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
        # 短暂注册一次回调
        ds = self.cam.data_stream[0]
        ds.register_capture_callback(one_time_cb)
        self.trig.send_command()
        # 等待回调写入 first_frame
        while self.first_frame is None:
            time.sleep(0.005)
        ds.unregister_capture_callback(one_time_cb)

    def stop_recording(self):
        self.recording = False
        if self.video_writer:
            self.video_writer.release()
        if self.log_file:
            self.log_file.close()

    def stop(self):
        self.running = False

# ----------------------- 辅助函数 -----------------------
def get_best_valid_bits(fmt):
    # 同原代码...
    return DxValidBit.BIT0_7

# ------------------------- GUI 界面 -------------------------
class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("相机采集与录制")
        self.worker = CameraWorker(cam_index=1)
        self.worker.frame_ready.connect(self.update_image)
        self.worker.start()

        # Image 显示
        self.label = QLabel()
        self.label.setFixedSize(640, 480)

        # 视频路径输入
        self.videoPathEdit = QLineEdit()
        btnVideo = QPushButton("选择视频保存路径")
        btnVideo.clicked.connect(self.select_video)

        # 时间戳路径输入
        self.logPathEdit = QLineEdit()
        btnLog = QPushButton("选择Log保存路径")
        btnLog.clicked.connect(self.select_log)

        # 开始/停止按钮
        self.btnStart = QPushButton("开始录制")
        self.btnStart.clicked.connect(self.toggle_record)

        # 布局
        h1 = QHBoxLayout()
        h1.addWidget(self.videoPathEdit)
        h1.addWidget(btnVideo)
        h2 = QHBoxLayout()
        h2.addWidget(self.logPathEdit)
        h2.addWidget(btnLog)
        h3 = QHBoxLayout()
        h3.addWidget(self.btnStart)

        v = QVBoxLayout(self)
        v.addWidget(self.label)
        v.addLayout(h1)
        v.addLayout(h2)
        v.addLayout(h3)

        # 定时触发采集
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.worker.trigger_once)
        self.timer.start(33)  # ≈30 FPS

        self.recording = False

    def select_video(self):
        path, _ = QFileDialog.getSaveFileName(self, "保存视频", "", "MP4 文件 (*.mp4)")
        if path:
            self.videoPathEdit.setText(path)

    def select_log(self):
        path, _ = QFileDialog.getSaveFileName(self, "保存时间戳", "", "文本文件 (*.txt)")
        if path:
            self.logPathEdit.setText(path)

    def toggle_record(self):
        if not self.recording:
            vp = self.videoPathEdit.text().strip()
            lp = self.logPathEdit.text().strip()
            if not vp or not lp:
                QMessageBox.warning(self, "警告", "请先选择视频和日志保存路径")
                return
            self.worker.start_recording(vp, lp)
            self.btnStart.setText("停止录制")
            self.recording = True
        else:
            self.worker.stop_recording()
            self.btnStart.setText("开始录制")
            self.recording = False

    def update_image(self, frame):
        """接收并显示来自 CameraWorker 的 BGR 帧"""
        h, w, _ = frame.shape
        img = QImage(frame.data, w, h, 3*w, QImage.Format_BGR888)
        self.label.setPixmap(QPixmap.fromImage(img))

    def closeEvent(self, e):
        """退出时清理线程"""
        self.worker.stop()
        self.worker.wait()
        super().closeEvent(e)

# ------------------------- 运行 -------------------------
if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())
